package com.googlecode.afdbaccess.web.common.servletfilter;
 
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Collection;

import javax.servlet.Filter;
import javax.servlet.FilterChain;
import javax.servlet.FilterConfig;
import javax.servlet.ServletException;
import javax.servlet.ServletOutputStream;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;

import org.apache.log4j.Logger;
import org.jboss.security.Base64Utils;
 
public class ETagFilter implements Filter {
    private final static Logger log = Logger.getLogger(ETagFilter.class);
    
    private boolean removeLastModifiedHeader = false;

    class DigestServletOutPutStreamWrapper extends ServletOutputStream {
        private ByteArrayOutputStream outputStream;
        private MessageDigest messageDigest = null;
        
        private ServletOutputStream servletOutputStream;
        
        public DigestServletOutPutStreamWrapper(ServletOutputStream servletOutputStream) throws NoSuchAlgorithmException {
            this.servletOutputStream = servletOutputStream;
            this.outputStream = new ByteArrayOutputStream();
            messageDigest = MessageDigest.getInstance("MD5");
        }
        
        @Override
        public void write(int value) throws IOException {
            outputStream.write(value);
            messageDigest.update((byte)value);
        }
        
        @Override
        public void write(byte[] b) throws IOException {
            outputStream.write(b);
            messageDigest.update(b);
        }
        
        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            outputStream.write(b, off, len);
            messageDigest.update(b, off, len);
        }
        
        public byte[] getDigest() {
            return messageDigest.digest();
        }
        
        @Override
        public void flush() throws IOException {
            servletOutputStream.write(outputStream.toByteArray());
            outputStream.reset();
        }

        public void reset() {
            outputStream.reset();
        }
    }
    
    class ResponseWrapper extends HttpServletResponseWrapper {
        private DigestServletOutPutStreamWrapper digestServletOutPutStreamWrapper;
        private PrintWriter printWriter;
        
        public ResponseWrapper(HttpServletResponse response) {
            super(response);
            //log.info("ResponseWrapper.constructor");
        }
        
        @Override
        public PrintWriter getWriter() throws IOException {
            //log.info("ResponseWrapper.getWriter()");
            if(printWriter == null) {
                String characterEncoding = getCharacterEncoding();
                //log.debug("characterEncoding: " + characterEncoding);
                if(characterEncoding == null || characterEncoding.trim().isEmpty()) {
                    characterEncoding = "iso-8859-1";
                }
                
                printWriter = new PrintWriter(new OutputStreamWriter(getOutputStream(), characterEncoding));
            }
            
            return printWriter;
        }
        
        @Override
        public ServletOutputStream getOutputStream() throws IOException {
            //log.info("ResponseWrapper.getOutputStream()");
            
            if(digestServletOutPutStreamWrapper == null) {
                try {
                    digestServletOutPutStreamWrapper = new DigestServletOutPutStreamWrapper(super.getOutputStream());
                } catch (NoSuchAlgorithmException nsae) {
                    throw new IOException("Couldn't find the digest", nsae);
                }
            }
            
            return digestServletOutPutStreamWrapper;
        }
        
        public byte[] getDigest() {
            //log.info("ResponseWrapper.getDigest()");
            
            if(digestServletOutPutStreamWrapper != null) {
                return digestServletOutPutStreamWrapper.getDigest();
            } else {
                return null;
            }
        }
        
        @Override
        public void flushBuffer() throws IOException {
            //log.info("ResponseWrapper.flushBuffer()");

            if(printWriter != null) {
                printWriter.flush();
            } else if(digestServletOutPutStreamWrapper != null) {
                digestServletOutPutStreamWrapper.flush();
            }
        }
        
        @Override
        public void resetBuffer() {
            if (digestServletOutPutStreamWrapper != null) {
                digestServletOutPutStreamWrapper.reset();
            }
            
            //super.resetBuffer();
        }
        
        @Override
        public void reset() {
            super.reset();
            resetBuffer();
        }
        
        @Override
        public void setHeader(String name, String value) {
            if(removeLastModifiedHeader && name.equals("Last-Modified")) {
                // do nothing. Swallow the header
            } else {
              super.setHeader(name, value);
            }
        }
    }
    
    @Override
    public void doFilter(ServletRequest req, ServletResponse res, FilterChain chain) throws IOException, ServletException {
        //log.info("doFilter()");
        HttpServletRequest request = (HttpServletRequest) req;
        
        if(request.getMethod().equals("GET") || request.getMethod().equals("HEAD")) {
            doEtagFilter(req, res, chain, request);
        } else {
            chain.doFilter(req, res);
        }
        
        //log.info("doFilter() end");
    }

    private void doEtagFilter(ServletRequest req, ServletResponse res, FilterChain chain, HttpServletRequest request) throws IOException, ServletException {
        //log.info("ETagFilter.doEtagFilter");
        
        HttpServletResponse response = (HttpServletResponse)res;

        ResponseWrapper responseWrapper = new ResponseWrapper(response);

        //log.info("doFilter() before call to doFilter");
        chain.doFilter(req, responseWrapper);
        //log.info("doFilter() after call to doFilter");

        String requestHeaderIfNoneMatch = request.getHeader("If-None-Match");

        if(requestHeaderIfNoneMatch != null || shouldGenerateETag(response)) {
            if(response.isCommitted()) {
                log.debug("Unable to modify ETag header since the response is already committed. status: " + response.getStatus());
            } else {
                String eTag = response.getHeader("ETag");

                if(eTag == null || eTag.startsWith("W/")) {
                    byte[] digest = responseWrapper.getDigest();
                    if(digest != null) {
                        eTag = "\"" + Base64Utils.tob64(digest) + "\"";
                        //log.debug("Generated ETag: " + eTag);
                    }
                }

                if (eTag != null && eTag.equals(requestHeaderIfNoneMatch)) {
                    clearUnneededHeaders(response);
                    //log.debug("Sending 304 since request ETag was: " + requestHeaderIfNoneMatch);
                    response.sendError(HttpServletResponse.SC_NOT_MODIFIED);
                }

                if(eTag != null) {
                    response.setHeader("ETag", eTag);
                }
            }
        }

        responseWrapper.flushBuffer();
    }

    public static void clearUnneededHeaders(HttpServletResponse response) {
        Collection<String> headerNames = response.getHeaderNames();
        
        for(String headerName : headerNames) {
            if(shouldRemoveHeader(headerName, response)) {
                response.setHeader(headerName, "");
            }
        }
    }
    
    public static boolean shouldRemoveHeader(String headerName, HttpServletResponse response) {
        if(headerName.equals("Expires") || headerName.endsWith("Set-Cookie")) {
            return false;
        }
        
        if(headerName.equals("Cache-Control")) {
            String cacheControlValue = response.getHeader("Cache-Control").replace(" ", "");
            
            if((cacheControlValue.indexOf("max-age=") >= 0 && cacheControlValue.indexOf("max-age=0") < 0) ||
               (cacheControlValue.indexOf("s-maxage=") >= 0 && cacheControlValue.indexOf("s-maxage=0") < 0)) {
                return false;
            }
        }
        
        return true;
    }

    private boolean shouldGenerateETag(HttpServletResponse response) {
        return
            (
                response.getHeader("Cache-Control") == null ||
                response.getHeader("Cache-Control").indexOf("no-store") < 0
            )
        ;
    }

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
        String removeLastModParam = filterConfig.getInitParameter("removeLastMod");
        if(removeLastModParam != null) {
            this.removeLastModifiedHeader = Boolean.parseBoolean(removeLastModParam);
        }
    }

    @Override
    public void destroy() {
        // do nothing
    }
}